import torch

from .super_token_association import ASSOCIATION_FNS
from .super_token_sampling import SAMPLING_FNS


def build_sampling_fn(cfg):
    cfg = cfg.copy() # avoid overwrite
    mode = cfg.pop("mode")
    def sampling_fn(token_positions, batch_token_ids):
        return SAMPLING_FNS[mode](
            token_positions, 
            batch_token_ids, 
            **cfg
        )
    return sampling_fn


def build_association_fn(cfg):
    cfg = cfg.copy() # avoid overwrite
    mode = cfg.pop("mode")
    def association_fn(token_positions, batch_token_ids, super_token_positions, batch_super_token_ids):
        return ASSOCIATION_FNS[mode](
            token_positions, 
            batch_token_ids, 
            super_token_positions,
            batch_super_token_ids,
            **cfg
        )
    return association_fn
        

class SuperTokenPreprocessor:
    """
    This preprocessor takes in voxel features and (optionally) image features\
    with their meta infomation and produces the super token initialization and\
    association indices for each stage.
    """
    def __init__(
        self,
        sampling_cfgs = [],
        association_cfgs = [],
    ) -> None:
        self.num_variants = len(sampling_cfgs)
        if isinstance(association_cfgs, dict):
            association_cfgs = [association_cfgs] * self.num_variants
        
        self.sampling_fns = [build_sampling_fn(cfg) for cfg in sampling_cfgs]
        self.association_fns = [build_association_fn(cfg) for cfg in association_cfgs]
        
    def __call__(
        self, 
        token_positions: torch.Tensor, 
        batch_token_ids
    ):
        """
        Returns a list of variants of super token information, each containing a dict:
        - super_token_src (torch.Tensor): [T] super_token index for each token
        - batch_super_token_ids (list[torch.Tensor]): super token indices for each batch
        - linked_pairs (torch.Tensor): [N x 2] N [token, super_token] index pairs (N > T)
        - num_linked_super_tokens (torch.Tensor): [T] number of linked super tokens
        - num_linked_tokens (torch.Tensor): [S] number of linked tokens
        """
        variants = []
        for v in range(self.num_variants):
            st_pos, st_src, batch_st_ids = self.sampling_fns[v](token_positions, batch_token_ids)
            linked_pairs, num_linked_super_tokens, num_linked_tokens = self.association_fns[v](
                token_positions, 
                batch_token_ids, 
                st_pos, 
                batch_st_ids
            )
            variants.append({
                "super_token_src": st_src,
                "super_token_positions": st_pos,
                "batch_super_token_ids": batch_st_ids,
                "linked_pairs": linked_pairs,
                "num_linked_super_tokens": num_linked_super_tokens,
                "num_linked_tokens": num_linked_tokens,
            })
        return variants